"""
    脚本文件说明：用于为已知onnx模型，添加预处理操作，比如，归一化、标准化等操作

"""

import onnx
import onnx.helper as helper
import numpy as np


def mergeOnnx():

    # 将目标模型和预处理模型合并到一起， 主要分为如下步骤：
    model = onnx.load("det.sim.onnx")
    pre_onnx = onnx.load("normal.onnx")


    # 1.修改预处理node和input、output名，防止与目标模型冲突
    for n in pre_onnx.graph.node:
        n.name = f"pre_{n.name}"
        
        for i in range(len(n.input)):
            n.input[i] = f"pre_{n.input[i]}"

        for i in range(len(n.output)):
            n.output[i] = f"pre_{n.output[i]}"

    # 2.目标模型第1个node的输入名称修改为预处理模型输出的名称，这样就算衔接上了
    for node in model.graph.node:
        if node.name == "p2o.Conv.0":
            node.input[0] = "pre_" + pre_onnx.graph.output[0].name
    # 插入节点，这种插入方式，顺序会乱
    # for n in pre_onnx.graph.node:
    #     model.graph.node.append(n)
    for index, n in enumerate(pre_onnx.graph.node):
        model.graph.node.insert(index, n)


    #第31、41行的名称为什么要修改？因为preonnx的节点输入输出都修改了，肯定也要对应修改输入输出的节点名称     
    model.graph.input[0].CopyFrom(pre_onnx.graph.input[0])
    model.graph.input[0].name = "pre_" +pre_onnx.graph.input[0].name
    # 当然可以二次修改input的名称，但一定也要修改对应承接该输入节点的输入名称
    # model.graph.input[0].name = "images"
    # model.graph.node[0].input[0] = "images"


    # 检查模型并打印
    onnx.checker.check_model(model)
    onnx.save(model, "det.sim.normal.onnx")


if __name__ == "__main__":
    mergeOnnx()



